﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

namespace Demo.Framework.Algorithms.UnionFind
{
    /*Aim - To determine whether two objects in anetwork are connected or not
     * We try to improve the Quick Union method by creating a balanced tree
     * Union(p,q) => 
     *  If p, q are isolated roots, put p under q
     *  Else if p is a part of a larger tree, ie.. Size(Root(p)) > Size(Root(q)), root(q) = root(p), vice versa
     *  IsConnected(p,q) => True if Root(p) = Root(q)
     * */
    class WeightedQuickUnion
    {
        private int[] _array;
        private int[] _size;
        private int _count = 0;

        public WeightedQuickUnion(int N)
        {
            _count = N;
            _array = new int[N];
            _size = new int[N];

            for (int i = 0; i < _count; i++)
            {
                _array[i] = i;
            }

            for (int i = 0; i < _count; i++)
            {
                _size[i] = 1;
            }
        }

        public bool IsConnected(int p, int q)
        {
            return Root(p) == Root(q);
        }

        /// <summary>
        /// Complexity - M * O(Log 2 N) , M = number of unions, N = number of array elements
        /// For each addition, the size of the tree can increase by a max of two
        /// So there are Log 2 N such operations possible
        /// </summary>
        /// <param name="p"></param>
        /// <param name="q"></param>
        /// <returns></returns>
        public bool Union(int p, int q)
        {
            try
            {
                int root_p = Root(p);
                int root_q = Root(q);

                if (_size[p] <= _size[q])
                {
                    // merge p into q
                    _array[root_p] = root_q;
                    _size[q] += 1;
                }
                else
                {
                    //merge q into p
                    _array[root_q] = root_p;
                    _size[p] += 1;
                }
                return true;
            }
            catch (Exception)
            {                
                return false;
            }
        }

        private int Root(int i)
        {
            while (i != _array[i])
            {
                i = _array[i];
            }
            return i;
        }

        public override string ToString()
        {
            StringBuilder sb = new StringBuilder();
            for (int i = 0; i < _count; i++)
            {
                if (i < _count - 1)
                {
                    sb.Append(_array[i] + ", ");
                }
                else
                {
                    sb.Append(_array[i]);
                }                
            }

            return sb.ToString();
        }
    }
}
